package com.hujiang.basic.search.common.utils;

import com.alibaba.fastjson.JSONObject;
import org.elasticsearch.search.aggregations.Aggregation;
import org.elasticsearch.search.aggregations.bucket.range.date.InternalDateRange;
import org.elasticsearch.search.aggregations.bucket.terms.InternalTerms;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.metrics.cardinality.InternalCardinality;
import org.joda.time.DateTime;

import java.util.*;


public class ElasticUtil {

    @SuppressWarnings({"rawtypes", "unchecked"})
    public static JSONObject formatAggs2JsonObject(Map<String, Aggregation> aggMap, int from) {
        JSONObject stat = new JSONObject();
        if (aggMap == null || aggMap.size() == 0) {
            return stat;
        }
        for (Map.Entry<String,Aggregation> entry : aggMap.entrySet()) {
            String key = entry.getKey();
            if (entry.getValue() instanceof InternalCardinality) {
                InternalCardinality agg = (InternalCardinality) (entry.getValue());
                long distinctAggCount = agg.getValue();
                stat.put(key, distinctAggCount);
            }
            if (entry.getValue() instanceof InternalTerms) {
                LinkedHashMap<String, Long> j = new LinkedHashMap<>();
                InternalTerms agg = (InternalTerms) (entry.getValue());
                List<Terms.Bucket> buckets = agg.getBuckets();
                if (!buckets.isEmpty() && from <= buckets.size()) {
                    for (int i = from; i < buckets.size(); i++) {
                        Terms.Bucket bucket = buckets.get(i);
                        j.put(bucket.getKey().toString(), bucket.getDocCount());
                    }
                    stat.put(key, j);
                }
            }
            if (entry.getValue() instanceof InternalDateRange) {
                InternalDateRange agg = (InternalDateRange) (entry.getValue());
                List<InternalDateRange.Bucket> buckets = agg.getBuckets();
                List<Map<String, Object>> ranges = new LinkedList<>();
                if (!buckets.isEmpty() && from <= buckets.size()) {
                    for (int i = from; i < buckets.size(); i++) {
                        Map<String, Object> j = new HashMap<>();
                        InternalDateRange.Bucket bucket = buckets.get(i);
                        j.put("doc_count", bucket.getDocCount());
                        if (null != bucket.getFrom()) {
                            j.put("from", ((DateTime) (bucket.getFrom())).getMillis());
                        }
                        if (null != bucket.getTo()) {
                            j.put("to", ((DateTime) (bucket.getTo())).getMillis());
                        }
                        ranges.add(j);
                    }
                    stat.put(key, ranges);
                }
            }
        }
        return stat;
    }
}
